# coding: UTF-8
import time
import torch
import argparse
import numpy as np
from importlib import import_module
from torch.utils.data import DataLoader

from utils.utils import get_time_dif
from dataset.CMN_dataset import CMNDataset
from train_eval.train_eval import train, test, init_network

if __name__ == '__main__':
    dataset = '.'
    
    model_name = 'DPCNN'

    x = import_module('models.' + model_name)
    config = (x.Config())
    np.random.seed(1)
    torch.manual_seed(1)
    torch.cuda.manual_seed_all(1)
    torch.backends.cudnn.deterministic = True
    
    start_time = time.time()
    print("Loading data...")
    train_dataset = CMNDataset(config, 'train')
    val_dataset = CMNDataset(config, 'val')
    test_dataset = CMNDataset(config, 'test')

    train_dataloader = DataLoader(
        train_dataset, 
        batch_size=config.batch_size, 
        shuffle=True,
        worker_init_fn=np.random.seed(1), 
    )
    val_dataloader = DataLoader(
        val_dataset, 
        batch_size=config.batch_size, 
        shuffle=True,
        worker_init_fn=np.random.seed(1), 
    )
    test_dataloader = DataLoader(
        test_dataset, 
        batch_size=config.batch_size, 
        shuffle=True,
        worker_init_fn=np.random.seed(1), 
    )
    time_dif = get_time_dif(start_time)
    print("Time usage:", time_dif)
    model = x.Model(config).to(config.device)
    init_network(model)
    print(model.parameters)
    train(
        config, 
        model, 
        train_dataloader, 
        val_dataloader, 
        test_dataloader, 
        # loss_function=FocalLoss(class_num=config.num_classes), 
    )

    '''
    model_names = ['Transformer', ]
    xs = []
    configs = []
    err = 0
    for model_name in model_names:
        x = import_module('models.' + model_name)
        xs.append(x)
        c = x.Config(dataset)
        configs.append(c)
        err = ((configs[0].device != c.device) or err)
        err = ((configs[0].batch_size != c.batch_size) or err)
        err = ((configs[0].pad_size != c.pad_size) or err)
    assert not(err)
    start_time = time.time()
    print("Loading data...")
    test_dataset = CMNDataset(configs, 'val')
    np.random.seed(1)
    torch.manual_seed(1)
    torch.cuda.manual_seed_all(1)
    torch.backends.cudnn.deterministic = True
    test_dataloader = DataLoader(
        test_dataset, 
        batch_size=configs[0].batch_size, 
        shuffle=True,
        worker_init_fn=np.random.seed(1), 
    )
    time_dif = get_time_dif(start_time)
    print("Time usage:", time_dif)
    models = []
    for i, x in enumerate(xs):
        model = x.Model(configs[i]).to(configs[0].device)
        models.append(model)
    test(
        configs, 
        models,  
        test_dataloader, 
        loss_function=FocalLoss(class_num=4), 
    )
    '''